# 读入tiff用的工具类
from osgeo import gdal
from torchvision.transforms import functional as F


def readTif_to_Ndarray(img_path):
    dataset = gdal.Open(img_path)
    # 获得矩阵的列数
    width = dataset.RasterXSize
    # 栅格矩阵的行数
    height = dataset.RasterYSize
    # 波段数
    bands = dataset.RasterCount
    # 获得数据
    data = dataset.ReadAsArray(0, 0, width, height)
    return data

class NdarrayToTensor:
    # torchvision.transforms.ToTensor()的重写
    def __init__(self):
        pass

    def __call__(self, pic):
        """
        Args:
            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.

        Returns:
            Tensor: Converted image.
        """
        return F.to_tensor(pic).permute(1, 0, 2)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"

def get_coordinates(tif_path):
    # 打开.tif文件
    dataset = gdal.Open(tif_path)

    # 获取地理变换信息
    transform = dataset.GetGeoTransform()

    # 获取左上角坐标
    top_left_x = transform[0]
    top_left_y = transform[3]

    # 获取像素大小
    pixel_width = transform[1]
    pixel_height = transform[5]

    # 获取图像的宽度和高度
    width = dataset.RasterXSize
    height = dataset.RasterYSize

    # 计算右下角坐标
    bottom_right_x = top_left_x + width * pixel_width
    bottom_right_y = top_left_y + height * pixel_height

    return {
        'top_left': (top_left_x, top_left_y),
        'bottom_right': (bottom_right_x, bottom_right_y),
        'width': width,
        'height': height
    }
